import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from model import LeNet



def test_data_process():
    test_data = FashionMNIST(root='./data',
                              train=False,
                              transform=transforms.Compose([transforms.Resize(size=28), transforms.ToTensor()]),
                              download=True)

    test_dataloader = Data.DataLoader(dataset=test_data,
                                       batch_size=1,
                                       shuffle=True,
                                       num_workers=0)
    return test_dataloader


def tester(model, test_dataloader):
    device = "cuda" if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
    test_corrects = 0.0
    test_num = 0

    model.eval()
    with torch.no_grad():
        for x, y in test_dataloader:
            x,y = x.to(device), y.to(device)
            output= model(x)
            pre_lab = torch.argmax(output, dim=1)
            test_corrects += torch.sum(pre_lab == y.data)
            test_num += x.size(0)

    # 计算测试准确率
    test_acc = test_corrects.double().item() / test_num
    print("测试的准确率为：", test_acc)




if __name__=="__main__":
    model = LeNet()
    model.load_state_dict(torch.load('./models/best_model.pth'))
    test_dataloader = test_data_process()
    tester(model, test_dataloader)